// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <melon/init/init.h>
#include <pollux/common/base/pollux_exception.h>

#include <pollux/expression/vector_function.h>
#include <pollux/functions/macros.h>

namespace kumo::pollux::functions {
    template<typename T>
    struct ArrayMinSimpleFunction {
        POLLUX_DEFINE_FUNCTION_TYPES(T);

        template<typename TInput>
        MELON_ALWAYS_INLINE bool call(
            TInput &out,
            const arg_type<Array<TInput> > &array) {
            if (array.size() == 0) {
                return false; // NULL
            }

            auto min = INT32_MAX;
            if (array.may_have_nulls()) {
                for (auto i = 0; i < array.size(); i++) {
                    if (!array[i].has_value()) {
                        return false; // NULL
                    }
                    auto v = array[i].value();
                    if (v < min) {
                        min = v;
                    }
                }
            } else {
                for (auto i = 0; i < array.size(); i++) {
                    auto v = array[i].value();
                    if (v < min) {
                        min = v;
                    }
                }
            }
            out = min;
            return true;
        }
    };

    template<typename T>
    struct ArrayMinSimpleFunctionIterator {
        POLLUX_DEFINE_FUNCTION_TYPES(T);

        template<typename TInput>
        MELON_ALWAYS_INLINE bool call(
            TInput &out,
            const arg_type<Array<TInput> > &array) {
            const auto size = array.size();
            if (size == 0) {
                return false; // NULL
            }

            auto min = INT32_MAX;
            if (array.may_have_nulls()) {
                for (const auto &item: array) {
                    if (!item.has_value()) {
                        return false; // NULL
                    }
                    auto v = item.value();
                    if (v < min) {
                        min = v;
                    }
                }
            } else {
                for (const auto &item: array) {
                    auto v = item.value();
                    if (v < min) {
                        min = v;
                    }
                }
            }

            out = min;
            return true;
        }
    };

    // Returns the minimum value in an array ignoring nulls.
    // The point of this is to exercise SkipNullsIterator.
    template<typename T>
    struct ArrayMinSimpleFunctionSkipNullIterator {
        POLLUX_DEFINE_FUNCTION_TYPES(T);

        template<typename TInput>
        MELON_ALWAYS_INLINE bool call(
            TInput &out,
            const arg_type<Array<TInput> > &array) {
            const auto size = array.size();
            if (size == 0) {
                return false; // NULL
            }

            bool hasValue = false;
            auto min = INT32_MAX;
            for (const auto &item: array.skipNulls()) {
                hasValue = true;
                if (item < min) {
                    min = item;
                }
            }

            if (!hasValue) {
                return false;
            }

            out = min;
            return true;
        }
    };

    // Basic vector function with out fast path.
    template<template <typename> class F, TypeKind kind>
    void applyTyped(
        const SelectivityVector &rows,
        const ArrayVector &array_vector,
        DecodedVector &elementsDecoded,
        VectorPtr &result) {
        using T = typename TypeTraits<kind>::NativeType;

        auto rawSizes = array_vector.rawSizes();
        auto rawOffsets = array_vector.rawOffsets();
        auto *flatResults = result->as_flat_vector<T>();
        rows.applyToSelected([&](auto row) {
            auto size = rawSizes[row];
            if (size == 0) {
                result->set_null(row, true);
                return;
            }

            auto offset = rawOffsets[row];
            auto vertex = elementsDecoded.value_at<T>(offset);

            for (auto i = offset; i < offset + size; i++) {
                if (elementsDecoded.is_null_at(i)) {
                    // If a NULL value is encountered, min/max are always NULL.
                    result->set_null(row, true);
                    return;
                }

                auto value = elementsDecoded.value_at<T>(i);
                if (F<T>()(value, vertex)) {
                    vertex = value;
                }
            }

            flatResults->set(row, vertex);
        });
    }

    // Decoder-based unoptimized implementation of min/max used to compare
    // performance of simple function min/max.
    template<template <typename> class F>
    class ArrayMinMaxFunctionBasic : public exec::VectorFunction {
    public:
        void apply(
            const SelectivityVector &rows,
            std::vector<VectorPtr> &args,
            const TypePtr & /* outputType */,
            exec::EvalCtx &context,
            VectorPtr &result) const override {
            POLLUX_CHECK_EQ(args.size(), 1);
            auto array_vector = args[0]->as_unchecked<ArrayVector>();

            auto elementsVector = array_vector->elements();
            exec::LocalSelectivityVector elementsRows(context, elementsVector->size());
            exec::LocalDecodedVector elementsHolder(
                context, *elementsVector, *elementsRows.get());

            BaseVector::ensure_writable(
                rows, elementsVector->type(), context.pool(), result);

            POLLUX_DYNAMIC_SCALAR_TEMPLATE_TYPE_DISPATCH(
                applyTyped,
                F,
                elementsVector->type_kind(),
                rows,
                *array_vector,
                *elementsHolder.get(),
                result);
        }
    };

    inline std::vector<std::shared_ptr<exec::FunctionSignature> > signatures() {
        static const std::vector<std::string> kSupportedTypeNames = {
            "boolean",
            "tinyint",
            "smallint",
            "integer",
            "bigint",
            "real",
            "double",
            "varchar",
            "timestamp",
        };

        std::vector<std::shared_ptr<exec::FunctionSignature> > signatures;
        signatures.reserve(kSupportedTypeNames.size());
        for (const auto &typeName: kSupportedTypeNames) {
            signatures.emplace_back(
                exec::FunctionSignatureBuilder()
                .returnType(typeName)
                .argumentType(fmt::format("array({})", typeName))
                .build());
        }
        return signatures;
    }

    POLLUX_DECLARE_VECTOR_FUNCTION(
        udf_array_min_basic,
        functions::signatures(),
        std::make_unique<functions::ArrayMinMaxFunctionBasic<std::less>>());
} // namespace kumo::pollux::functions
